package utils

import (
	"crypto"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha256"
	"crypto/x509"
	"encoding/pem"
	"errors"
	"io/ioutil"
	"os"
)

// GenerateRSAKey
// @Author: gh_shidawu
// @Date: 2022-01-13 17:10:37
// @Description: 生成RSA私钥和公钥，保存到文件中
func GenerateRSAKey(bits int, privateKeyFile, publicKeyFile string) error {
	// GenerateKey函数使用随机数据生成器random生成一对具有指定字位数的RSA密钥
	// Reader是一个全局、共享的密码用强随机数生成器
	privateKey, err := rsa.GenerateKey(rand.Reader, bits)
	if err != nil {
		return err
	}

	// 通过x509标准将得到的ras私钥序列化为ASN.1 的 DER编码字符串
	X509PrivateKey := x509.MarshalPKCS1PrivateKey(privateKey)

	// 创建私钥文件
	privateFile, err := os.Create(privateKeyFile)
	if err != nil {
		return err
	}
	defer privateFile.Close()

	// 使用pem格式对x509输出的内容进行编码
	// 构建一个pem.Block结构体对象
	privateBlock := pem.Block{Type: "RSA Private Key", Bytes: X509PrivateKey}

	// 将私钥数据保存到文件
	err = pem.Encode(privateFile, &privateBlock)

	// 获取公钥的数据
	publicKey := privateKey.PublicKey

	// X509对公钥编码
	X509PublicKey, err := x509.MarshalPKIXPublicKey(&publicKey)
	if err != nil {
		return err
	}

	// 创建公钥文件
	publicFile, err := os.Create(publicKeyFile)
	if err != nil {
		return err
	}
	defer publicFile.Close()

	// pem格式编码
	// 创建一个pem.Block结构体对象
	publicBlock := pem.Block{Type: "RSA Public Key", Bytes: X509PublicKey}

	// 将公钥数据保存到文件
	err = pem.Encode(publicFile, &publicBlock)

	return err
}

// ReadRSAKey
// @Author: lee
// @Date: 2022-01-13 16:39:14
// @Description: 读取RSA文件内容
func ReadRSAKey(keyFile string) (content []byte, err error) {
	// 读取公钥内容
	file, err := os.Open(keyFile)
	if err != nil {
		return nil, err
	}
	defer file.Close()

	return ioutil.ReadAll(file)
}

// GetRSAPublicKey
// @Author: lee
// @Date: 2022-01-13 16:38:59
// @Description: 获取RSA公钥结构体
func GetRSAPublicKey(content []byte) (publicKey *rsa.PublicKey, err error) {
	// pem解码
	block, _ := pem.Decode(content)
	if block == nil {
		return nil, errors.New("public key error")
	}

	// x509解码
	publicKeyInterface, err := x509.ParsePKIXPublicKey(block.Bytes)
	if err != nil {
		return nil, err
	}

	// 类型断言
	publicKey = publicKeyInterface.(*rsa.PublicKey)
	return publicKey, nil
}

// @Author: lee
// @Date: 2022-01-13 16:38:32
// @Description: 获取RSA私钥结构体
func GetRSAPrivateKey(content []byte) (privateKey *rsa.PrivateKey, err error) {
	//pem解码
	block, _ := pem.Decode(content)
	if block == nil {
		return nil, errors.New("private key error!")
	}

	//X509解码
	privateKey, err = x509.ParsePKCS1PrivateKey(block.Bytes)
	return privateKey, nil
}

// RSAEncrypt
// @Author: lee
// @Date: 2022-01-13 15:11:31
// @Description: RSA加密
func RSAEncrypt(plaintext, publicKey []byte) (ciphertext []byte, err error) {
	// 获取公钥结构体
	pub, err := GetRSAPublicKey(publicKey)
	if err != nil {
		return nil, err
	}

	//加密
	return rsa.EncryptPKCS1v15(rand.Reader, pub, plaintext)
}

// RSADecrypt
// @Author: lee
// @Date: 2022-01-13 15:11:31
// @Description: RSA解密
func RSADecrypt(ciphertext, privateKey []byte) (plaintext []byte, err error) {
	// 获取私钥结构体
	priv, err := GetRSAPrivateKey(privateKey)
	if err != nil {
		return nil, err
	}

	// 解密
	return rsa.DecryptPKCS1v15(rand.Reader, priv, ciphertext)
}

// SignRSA
// @Author: lee
// @Date: 2022-01-13 16:38:06
// @Description: RSA签名
func SignRSA(plaintext, privateKey []byte) (sign []byte, err error) {
	// 获取私钥结构体
	priv, err := GetRSAPrivateKey(privateKey)
	if err != nil {
		return nil, err
	}

	// 计算散列值
	hash := sha256.New()
	hash.Write(plaintext)
	bytes := hash.Sum(nil)

	// 签名
	return rsa.SignPKCS1v15(rand.Reader, priv, crypto.SHA256, bytes)
}

// VerifyRSA
// @Author: lee
// @Date: 2022-01-13 16:37:28
// @Description: 验证RSA签名
func VerifyRSA(plaintext, sign, publicKey []byte) (verified bool, err error) {
	// 获取公钥结构体
	pub, err := GetRSAPublicKey(publicKey)
	if err != nil {
		return false, err
	}

	// 计算散列值
	hash := sha256.New()
	hash.Write(plaintext)
	bytes := hash.Sum(nil)

	// 验证签名
	err = rsa.VerifyPKCS1v15(pub, crypto.SHA256, bytes, sign)

	return err == nil, err
}
